/*
 * Copyright 2025 Snowflake Inc.
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "postgres.h"

#include "access/tupdesc.h"
#include "commands/copy.h"
#include "commands/defrem.h"
#include "pg_lake/cleanup/in_progress_files.h"
#include "pg_lake/copy/copy_format.h"
#include "pg_lake/csv/csv_options.h"
#include "pg_lake/csv/csv_writer.h"
#include "pg_lake/fdw/data_files_catalog.h"
#include "pg_lake/fdw/multi_data_file_dest.h"
#include "pg_lake/fdw/writable_table.h"
#include "pg_lake/fdw/schema_operations/register_field_ids.h"
#include "pg_lake/pgduck/remote_storage.h"
#include "pg_lake/pgduck/write_data.h"
#include "pg_lake/storage/local_storage.h"
#include "foreign/foreign.h"
#include "tcop/dest.h"
#include "utils/memutils.h"
#include "utils/palloc.h"
#include "utils/rel.h"


/* pg_lake_table.max_write_temp_file_size_mb setting */
int			MaxWriteTempFileSizeMB = DEFAULT_MAX_WRITE_TEMP_FILE_SIZE_MB;


/*
 * MultiDataFileUploadDestReceiver is a DestReceiver that writes the tuples
 * to a chain of files generated using an internal DestReceiver.
 *
 * When the file size exceeds MaxWriteTempFileSizeMB, the current file
 * is converted to a data file and a new one is started. Data files are
 * stored in S3 and can therefore be removed from the local machine once
 * written.
 */
typedef struct MultiDataFileUploadDestReceiver
{
	/* function pointers for this DestReceiver */
	DestReceiver pub;

	/* function pointers for the active CSV DestReceiver */
	DestReceiver *currentDest;

	/* current CSV file */
	char	   *currentFilePath;

	/* current number of rows in the CSV file */
	int64		currentRowCount;

	/* starting row id - if non-zero, files added will start with that rowid */
	/* see CreateMultiDataFileDestReceiver for more details */
	int64		currentRowIdStart;

	/* relation to which we are writing */
	Oid			relationId;

	/* memory context for the DestReceiver */
	MemoryContext parentContext;

	/* memory context for the child DestReceiver */
	MemoryContext childContext;

	/* target format of the DestReceiver */
	CopyDataFormat targetFormat;

	/* schema of the table */
	DataFileSchema *schema;

	/* tuple descriptor of the received slots */
	TupleDesc	tupleDesc;

	/* operation of the DestReceiver */
	int			operation;

	/* modifications generated by this DestReceiver */
	List	   *modifications;

	/* the relation's partition spec id */
	int32		currentPartitionSpecId;

	/*
	 * Some callers might prefer to use a reduced value (e.g., partitioned
	 * dest receiver).
	 */
	int			maxWriteTempFileSizeMB;
}			MultiDataFileUploadDestReceiver;


static void CreateChildDestReceiver(MultiDataFileUploadDestReceiver * self);
static void FlushChildDestReceiver(MultiDataFileUploadDestReceiver * self);
static void StartDestReceiver(DestReceiver *self, int operation, TupleDesc typeinfo);
static bool ReceiveSlot(TupleTableSlot *slot, DestReceiver *self);
static void ShutdownDestReceiver(DestReceiver *self);
static void DestroyDestReceiver(DestReceiver *self);

/*
 * CreateMultiDataFileDestReceiver creates a
 * MultiDataFileUploadDestReceiver with the given target format and file
 * size limit. In other words, the splitting is based on the file size.
 *
 * When rowIdStart is 0, rowId assignment is automatically handled by the
 * DestReceiver (specifically PrepareCSVInsertion()).  This is the common case
 * for inserts or imports of data files into a table.  This will ensure that
 * the rowmap records are properly sync'd with the file ids that are generated
 * at the time, and will be handled automatically.
 *
 * The internal handling of rowids relies on being able to map the given
 * Relation to the underlying rowid sequence and ensure that we can map a
 * specific range of rowids to rows in a file.  This means that we need to
 * lock the sequence relation used to store the starting row number until we
 * know how many rows need to be allocated.  (See row_ids.c for more grotty
 * details.)
 *
 * When rowIdStart is non-zero, the caller takes on the responsibility of
 * ensuring that the sequence is appropriately locked to ensure that no
 * concurrent callers can use a rowId that will later be assigned.  This is
 * currently only used by the logical replication copy_data handler to ensure
 * that the row ids are pre-reserved and we can avoid needing to update the
 * replay tables with the actual discovered rowids.
 *
 * StartReservingRowIdRange() and FinishReservingRowIdRange() should likely be
 * used for any future callers as well.
 */
DestReceiver *
CreateMultiDataFileDestReceiver(Oid relationId,
								CopyDataFormat targetFormat,
								int MaxWriteTempFileSizeMB,
								int32 partitionSpecId,
								int64 rowIdStart)
{
	MultiDataFileUploadDestReceiver *self =
		(MultiDataFileUploadDestReceiver *) palloc0(sizeof(MultiDataFileUploadDestReceiver));

	self->pub.rStartup = StartDestReceiver;
	self->pub.receiveSlot = ReceiveSlot;
	self->pub.rShutdown = ShutdownDestReceiver;
	self->pub.rDestroy = DestroyDestReceiver;
	self->pub.mydest = DestCopyOut;

	self->relationId = relationId;
	self->maxWriteTempFileSizeMB = MaxWriteTempFileSizeMB;
	self->targetFormat = targetFormat;
	self->parentContext = CurrentMemoryContext;
	self->childContext = AllocSetContextCreate(CurrentMemoryContext,
											   "MultiDataFileUploadDestReceiver",
											   ALLOCSET_DEFAULT_MINSIZE,
											   ALLOCSET_DEFAULT_INITSIZE,
											   ALLOCSET_DEFAULT_MAXSIZE);
	self->currentRowIdStart = rowIdStart;

	/* cache the schema field */
	self->schema = GetDataFileSchemaForTable(relationId);
	self->currentPartitionSpecId = partitionSpecId;

	return (DestReceiver *) self;
}


/*
 * CreateChildDestReceiver creates the DestReceiver for the current file.
 */
static void
CreateChildDestReceiver(MultiDataFileUploadDestReceiver * self)
{
	/* use the same memory context that was used to create the DestReceiver */
	MemoryContext oldContext = MemoryContextSwitchTo(self->childContext);

	/* we currently use CSV as a universal intermediate format */
	bool		includeHeader = true;
	List	   *copyOptions = InternalCSVOptions(includeHeader);
	char	   *tempFilePath = GenerateTempFileName("lake_table_insert", true);

	self->currentRowCount = 0;
	self->currentFilePath = tempFilePath;
	self->currentDest = CreateCSVDestReceiver(tempFilePath, copyOptions,
											  self->targetFormat);

	self->currentDest->rStartup(self->currentDest, self->operation, self->tupleDesc);

	MemoryContextSwitchTo(oldContext);
}


/*
 * FlushChildDestReceiver finalizes the write file and converts it to
 * one or more data files.
 */
static void
FlushChildDestReceiver(MultiDataFileUploadDestReceiver * self)
{
	self->currentDest->rShutdown(self->currentDest);

	/* do conversion in a memory context that is about to be destroyed */
	MemoryContext oldContext = MemoryContextSwitchTo(self->childContext);

	List	   *insertModifications =
		PrepareCSVInsertion(self->relationId,
							self->currentFilePath,
							self->currentRowCount,
							self->currentRowIdStart,
							GetCSVDestReceiverMaxLineSize(self->currentDest),
							self->schema);

	/* make sure we preserve the list of data file modifications */
	MemoryContextSwitchTo(self->parentContext);
	ListCell   *modificationCell = NULL;

	foreach(modificationCell, insertModifications)
	{
		DataFileModification *modification = lfirst(modificationCell);
		DataFileModification *copyModification = palloc0(sizeof(DataFileModification));

		copyModification->type = ADD_DATA_FILE;
		copyModification->insertFile = pstrdup(modification->insertFile);
		copyModification->insertedRowCount = modification->insertedRowCount;

		copyModification->partitionSpecId = self->currentPartitionSpecId;
		copyModification->partition = modification->partition;

		/*
		 * If caller of dest receiver is assigning rowids itself,
		 * currentRowIdStart will be > 0, otherwise 0 if totally managed.
		 *
		 * When initially assigned by the caller, we need to increase the
		 * value ourselves here to match the underlying row counts for the
		 * files generated by PrepareCSVInsertion().
		 */
		if (self->currentRowIdStart > 0)
		{
			copyModification->reservedRowIdStart = self->currentRowIdStart;

			/* increment currentRowIdStart by number of rows */
			self->currentRowIdStart += modification->insertedRowCount;
		}

		/* add a data mapping for each file  */
		self->modifications = lappend(self->modifications, copyModification);
	}

	MemoryContextSwitchTo(oldContext);

	self->currentDest->rDestroy(self->currentDest);
	self->currentDest = NULL;

	/* release all memory allocated for child DestReceiver, and the temp file */
	MemoryContextReset(self->childContext);
}


/*
 * StartDestReceiver is called when the DestReceiver starts.
 */
static void
StartDestReceiver(DestReceiver *dest, int operation, TupleDesc tupleDesc)
{
	MultiDataFileUploadDestReceiver *self = (MultiDataFileUploadDestReceiver *) dest;

	self->operation = operation;
	self->tupleDesc = tupleDesc;

	CreateChildDestReceiver(self);
}


/*
 * ReceiveSlot is called when a slot is received. We pass it on to the child
 * DestReceiver, and rotate to a new file if the file size exceeds the threshold.
 */
static bool
ReceiveSlot(TupleTableSlot *slot, DestReceiver *dest)
{
	MultiDataFileUploadDestReceiver *self = (MultiDataFileUploadDestReceiver *) dest;

	if (self->currentDest == NULL)
		CreateChildDestReceiver(self);

	bool		result = self->currentDest->receiveSlot(slot, self->currentDest);

	self->currentRowCount++;

	if (GetCSVDestReceiverFileSize(self->currentDest) > self->maxWriteTempFileSizeMB * MB_BYTES)
		FlushChildDestReceiver(self);

	return result;
}


/*
 * ShutdownDestReceiver is called at the end of a write. We finalize the last
 * write file and convert it to one or more data files.
 */
static void
ShutdownDestReceiver(DestReceiver *dest)
{
	MultiDataFileUploadDestReceiver *self = (MultiDataFileUploadDestReceiver *) dest;

	if (self->currentDest != NULL)
		FlushChildDestReceiver(self);
}


/*
 * DestroyDestReceiver destroys the DestReceiver. We currently don't do any
 * special cleanup steps, since the relevant work already happens in shutdown.
 */
static void
DestroyDestReceiver(DestReceiver *dest)
{
}



/*
 * GetMultiDataFileDestReceiverModifications returns the modifications generated
 * by this DestReceiver.
 */
List *
GetMultiDataFileDestReceiverModifications(DestReceiver *dest)
{
	MultiDataFileUploadDestReceiver *self = (MultiDataFileUploadDestReceiver *) dest;

	return self->modifications;
}
